#include <strmif.h>
#include "mkvparserstreamaudio.hpp"
#include "mkvparser.hpp"
#include "fourccguid.hpp"
#include "vorbistypes.hpp"
#include <cassert>
#include <limits>
#include <uuids.h>
#include <mmreg.h>
#include <malloc.h>
#ifdef _DEBUG
#include "odbgstream.hpp"
using std::endl;
#endif


namespace MkvParser
{

AudioStream* AudioStream::CreateInstance(AudioTrack* pTrack)
{
    assert(pTrack);
    
    const char* const id = pTrack->GetCodecId();
    assert(id);  //TODO
    
    if (_stricmp(id, "A_AC3") == 0)
        __noop;
    else if (_stricmp(id, "A_VORBIS") == 0)
        __noop;
    else if (_stricmp(id, "A_MPEG/L3") == 0)
        __noop;
    else
        return 0;  //can't create a stream from this track
        
    //TODO: vet settings, etc    

    AudioStream* const s = new (std::nothrow) AudioStream(pTrack);
    assert(s);  //TODO
    
    return s;
}


AudioStream::AudioStream(AudioTrack* pTrack) :
    Stream(pTrack),
    m_preroll(0)
{
}


std::wostream& AudioStream::GetKind(std::wostream& os) const
{
    return os << L"Audio";
}


BYTE AudioStream::GetChannels() const
{
    const AudioTrack* const pTrack = static_cast<AudioTrack*>(m_pTrack);
    
    const __int64 channels = pTrack->GetChannels();
    assert(channels > 0);
    assert(channels <= 255);
    
    const BYTE result = static_cast<BYTE>(channels);
    return result;
}


ULONG AudioStream::GetSamplesPerSec() const
{
    const AudioTrack* const pTrack = static_cast<AudioTrack*>(m_pTrack);

    const double rate = pTrack->GetSamplingRate();
    assert(rate > 0);
    
    double intrate_;
    const double fracrate = modf(rate, &intrate_);
    fracrate;
    assert(fracrate == 0);
    
    const ULONG result = static_cast<ULONG>(intrate_);
    return result;
}


BYTE AudioStream::GetBitsPerSample() const
{
    const AudioTrack* const pTrack = static_cast<AudioTrack*>(m_pTrack);
    
    const __int64 val = pTrack->GetBitDepth();
    
    if (val <= 0)
        return 0;
        
    assert(val < 256);
    assert((val % 8) == 0);
    
    const BYTE result = static_cast<BYTE>(val);
    return result;
}


void AudioStream::GetMediaTypes(CMediaTypes& mtv) const
{
    mtv.Clear();
    
    const char* const id = m_pTrack->GetCodecId();
    assert(id);
    
    if (_stricmp(id, "A_AC3") == 0)
        GetAc3MediaTypes(mtv);
        
    else if (_stricmp(id, "A_VORBIS") == 0)
        GetVorbisMediaTypes(mtv);

    else if (_stricmp(id, "A_MPEG/L3") == 0)
        GetMp3MediaTypes(mtv);
}


void AudioStream::GetAc3MediaTypes(CMediaTypes& mtv) const
{
    AM_MEDIA_TYPE mt;
    WAVEFORMATEX wfx;
    
    mt.majortype = MEDIATYPE_Audio;
    mt.subtype = MEDIASUBTYPE_DOLBY_AC3;
    mt.bFixedSizeSamples = FALSE;
    mt.bTemporalCompression = FALSE;
    mt.lSampleSize = 0;
    mt.formattype = FORMAT_WaveFormatEx;
    mt.pUnk = 0;
    mt.cbFormat = sizeof wfx;
    mt.pbFormat = (BYTE*)&wfx;
    
    wfx.wFormatTag = WAVE_FORMAT_UNKNOWN;  //TODO: see codec.pdf (at matroska website)
    wfx.nChannels = GetChannels();
    wfx.nSamplesPerSec = GetSamplesPerSec();
    wfx.nAvgBytesPerSec = 0;  //TODO: use s_frmsizcod_table for this
    wfx.nBlockAlign = 0;      //TODO
    wfx.wBitsPerSample = 0;   //TODO
    wfx.cbSize = 0;
    
    mtv.Add(mt);    
}


void AudioStream::GetVorbisMediaTypes(CMediaTypes& mtv) const
{
    const bytes_t& cp = m_pTrack->GetCodecPrivate();
    assert(!cp.empty());
    
    const ULONG cp_size = static_cast<ULONG>(cp.size());
    assert(cp_size > 0);
    
    const BYTE* const begin = &cp[0];
    const BYTE* const end = begin + cp_size;
    
    const BYTE* p = begin;
    assert(p < end);
    
    const BYTE n = *p++;
    n;
    assert(n == 2);
    assert(p < end);
    
    const ULONG id_len = *p++;  //TODO: don't assume < 255
    assert(id_len < 255);
    assert(id_len > 0);
    assert(p < end);
    
    const ULONG comment_len = *p++;  //TODO: don't assume < 255
    assert(comment_len < 255);
    assert(comment_len > 0);
    assert(p < end);
    
    //p points to first header
    
    const BYTE* const id_hdr = p;
    id_hdr;
    
    const BYTE* const comment_hdr = id_hdr + id_len;
    comment_hdr;
    
    const BYTE* const setup_hdr = comment_hdr + comment_len;
    setup_hdr;
    assert(setup_hdr < end);
    
    const ptrdiff_t setup_len_ = end - setup_hdr;
    assert(setup_len_ > 0);
    
    const DWORD setup_len = static_cast<DWORD>(setup_len_);
    
    const size_t hdr_len = id_len + comment_len + setup_len;

    using VorbisTypes::VORBISFORMAT2;
    
    const size_t cb = sizeof(VORBISFORMAT2) + hdr_len;
    BYTE* const pb = (BYTE*)_malloca(cb);
    
    VORBISFORMAT2& fmt = (VORBISFORMAT2&)(*pb);
    
    fmt.dwChannels = GetChannels();
    fmt.dwSamplesPerSec = GetSamplesPerSec();
    fmt.dwBitsPerSample = GetBitsPerSample();
    fmt.dwHeaderSizes[0] = id_len;
    fmt.dwHeaderSizes[1] = comment_len;
    fmt.dwHeaderSizes[2] = setup_len;
    
    assert(p < end);
    assert(size_t(end - p) == hdr_len);

    BYTE* const dst = pb + sizeof(VORBISFORMAT2);
    memcpy(dst, p, hdr_len);
    
    AM_MEDIA_TYPE mt;
    
    mt.majortype = MEDIATYPE_Audio;
    mt.subtype = VorbisTypes::MEDIASUBTYPE_VORBIS2_AUDIO;
    mt.bFixedSizeSamples = FALSE;
    mt.bTemporalCompression = FALSE;
    mt.lSampleSize = 0;
    mt.formattype = VorbisTypes::FORMAT_VorbisFormat2;
    mt.pUnk = 0;
    mt.cbFormat = static_cast<ULONG>(cb);
    mt.pbFormat = pb;
    
    mtv.Add(mt);
    
    mt.majortype = VorbisTypes::MEDIATYPE_OggPacketStream;
    mt.subtype = MEDIASUBTYPE_None;
    mt.bFixedSizeSamples = FALSE;
    mt.bTemporalCompression = FALSE;
    mt.lSampleSize = 0;
    mt.formattype = VorbisTypes::FORMAT_OggIdentHeader;
    mt.pUnk = 0;
    mt.cbFormat = id_len;
    mt.pbFormat = const_cast<BYTE*>(id_hdr);
    
    mtv.Add(mt);
}


void AudioStream::GetMp3MediaTypes(CMediaTypes& mtv) const
{
    AM_MEDIA_TYPE mt;
    
    MPEGLAYER3WAVEFORMAT wfmp3;
    WAVEFORMATEX& wfx = wfmp3.wfx;
    
    mt.majortype = MEDIATYPE_Audio;
    mt.subtype = MJH::FourCCGUID(WAVE_FORMAT_MPEGLAYER3);
    mt.bFixedSizeSamples = TRUE;
    mt.bTemporalCompression = FALSE;
    mt.lSampleSize = 1;
    mt.formattype = FORMAT_WaveFormatEx;
    mt.pUnk = 0;
    mt.cbFormat = sizeof wfmp3;
    mt.pbFormat = (BYTE*)&wfmp3;
    
    wfx.wFormatTag = WAVE_FORMAT_MPEGLAYER3;
    wfx.nChannels = GetChannels();
    wfx.nSamplesPerSec = GetSamplesPerSec();
    wfx.nAvgBytesPerSec = 0;  //TODO
    wfx.nBlockAlign = 1;
    wfx.wBitsPerSample = GetBitsPerSample();
    wfx.cbSize = MPEGLAYER3_WFX_EXTRA_BYTES;
    
    wfmp3.wID = 0;  //TODO
    wfmp3.fdwFlags = 0;  //TODO
    wfmp3.nBlockSize = 0;  //TODO
    wfmp3.nFramesPerBlock = 0;  //TODO
    wfmp3.nCodecDelay = 0;  //TODO
    
    mtv.Add(mt);    
}


HRESULT AudioStream::QueryAccept(const AM_MEDIA_TYPE* pmt) const
{
    if (pmt == 0)
        return E_INVALIDARG;
        
    const AM_MEDIA_TYPE& mt = *pmt;
    
    if (mt.majortype != MEDIATYPE_Audio)
        return S_FALSE;
        
    const char* const id = m_pTrack->GetCodecId();
    assert(id);
    
    if (_stricmp(id, "A_AC3") == 0)
    {
        if (mt.subtype != MEDIASUBTYPE_DOLBY_AC3)
            return S_FALSE;
            
        return S_OK;  //TODO: more vetting here
    }
    else if (_stricmp(id, "A_VORBIS") == 0)
    {
        if (mt.subtype != VorbisTypes::MEDIASUBTYPE_VORBIS2_AUDIO)
            return S_FALSE;
            
        return S_OK;
    }
    else if (_stricmp(id, "A_MPEG/L3") == 0)
    {
        if (!MJH::FourCCGUID::IsFourCC(mt.subtype))
            return S_FALSE;
            
        const DWORD dwFormatTag = mt.subtype.Data1;
        
        if (dwFormatTag != WAVE_FORMAT_MPEGLAYER3)
            return S_FALSE;
            
        return S_OK;
    }

    return S_FALSE;  
}


HRESULT AudioStream::SetConnectionMediaType(const AM_MEDIA_TYPE& mt)
{
    if (mt.majortype == VorbisTypes::MEDIATYPE_OggPacketStream)
        m_preroll = &AudioStream::SendOggIdentPacket;
    else
        m_preroll = &AudioStream::DoNothing;

    return S_OK;
}


HRESULT AudioStream::UpdateAllocatorProperties(
    ALLOCATOR_PROPERTIES& props) const
{
    if (props.cBuffers <= 0)
        props.cBuffers = 1;
        
    const long size = GetBufferSize();
    
    if (props.cbBuffer < size)
        props.cbBuffer = size;
        
    if (props.cbAlign <= 0)
        props.cbAlign = 1;
        
    if (props.cbPrefix < 0)
        props.cbPrefix = 0;
    
    return S_OK;
}


long AudioStream::GetBufferSize() const
{
    const AudioTrack* const pTrack = static_cast<AudioTrack*>(m_pTrack);
    
    const double rr = pTrack->GetSamplingRate();
    const long nSamplesPerSec = static_cast<long>(rr);
    
    const __int64 cc = pTrack->GetChannels();
    const long nChannels = static_cast<long>(cc);
    
    const long nBlockAlign = nChannels * 2;  //16-bit PCM
    const long nAvgBytesPerSec = nBlockAlign * nSamplesPerSec;
        
    const long size = nAvgBytesPerSec;  //TODO: make a better estimate
    
    return size;
}


HRESULT AudioStream::OnPopulateSample(
    const BlockEntry* pNextEntry,
    IMediaSample* pSample)
{
    assert(pSample);
    assert(m_pBase);
    assert(!m_pBase->EOS());
    assert(m_pCurr);
    assert(m_pCurr != m_pStop);
    assert(!m_pCurr->EOS());
    
    const Block* const pCurrBlock = m_pCurr->GetBlock();
    assert(pCurrBlock);    
    assert(pCurrBlock->GetNumber() == m_pTrack->GetNumber());
    
    Cluster* const pCurrCluster = m_pCurr->GetCluster();
    assert(pCurrCluster);
    
    assert((m_pStop == 0) || 
           m_pStop->EOS() ||
           (m_pStop->GetBlock()->GetTimeCode(m_pStop->GetCluster()) > 
             pCurrBlock->GetTimeCode(pCurrCluster)));
    
    const __int64 basetime_ns = m_pBase->GetTime();
    assert(basetime_ns >= 0);
    assert((basetime_ns % 100) == 0);
    assert(pCurrBlock->GetTime(pCurrCluster) >= basetime_ns);
    
    const LONG srcsize = pCurrBlock->GetSize();
    assert(srcsize >= 0);
    
    const long tgtsize = pSample->GetSize();
    tgtsize;
    assert(tgtsize >= 0);
    assert(tgtsize >= srcsize);
    
    IAsyncReader* const pReader = pCurrCluster->m_pSegment->m_pReader;
    
    BYTE* ptr;

    HRESULT hr = pSample->GetPointer(&ptr);
    assert(SUCCEEDED(hr));
    assert(ptr);
    
    hr = pCurrBlock->Read(pReader, ptr);
    assert(hr == S_OK);  //all bytes were read
    
    hr = pSample->SetActualDataLength(srcsize);
    
    hr = pSample->SetPreroll(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetMediaType(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetDiscontinuity(m_bDiscontinuity);
    assert(SUCCEEDED(hr));
    
    //set by caller:
    //m_bDiscontinuity = false;

    //TODO: this is meaningful for audio streams.  MediaTime is the accumulated
    //sum of audio samples.
    hr = pSample->SetMediaTime(0, 0);
    assert(SUCCEEDED(hr));
    
    //const BlockGroup* const pGroup = m_pCurr->m_pGroup;
    //pGroup;
    //assert(pGroup);
    
    //TODO: you could make an argument that every audio frame is a key frame.
    //Here we decide by fiat that discussion of sync point is only meaningful
    //for video streams, so for audio sync point status is always false.
    hr = pSample->SetSyncPoint(FALSE);
    assert(SUCCEEDED(hr));
    
    const __int64 start_ns = pCurrBlock->GetTime(pCurrCluster);
    assert(start_ns >= basetime_ns);
    assert((start_ns % 100) == 0);
    
    __int64 ns = start_ns - basetime_ns;    
    __int64 start_reftime = ns / 100;
    
    __int64 stop_reftime;
    __int64* pstop_reftime;
            
    if ((pNextEntry == 0) || pNextEntry->EOS())
        pstop_reftime = 0;  //TODO: use duration of curr block
    else
    {
        const Block* const pNextBlock = pNextEntry->GetBlock();
        assert(pNextBlock);
        
        Cluster* const pNextCluster = pNextEntry->GetCluster();
        
        const __int64 stop_ns = pNextBlock->GetTime(pNextCluster);
        assert(stop_ns > start_ns);
        assert((stop_ns % 100) == 0);

        ns = stop_ns - basetime_ns;
        stop_reftime = ns / 100;
        assert(stop_reftime > start_reftime);

        pstop_reftime = &stop_reftime;
    }
    
    hr = pSample->SetTime(&start_reftime, pstop_reftime);
    assert(SUCCEEDED(hr));
    
    //set by caller:
    //m_pCurr = pNextBlock;    
    return S_OK;
}


bool AudioStream::SendPreroll(IMediaSample* pSample)
{
    assert(m_preroll);
    return (this->*m_preroll)(pSample);
}    


bool AudioStream::DoNothing(IMediaSample*)
{
    return false;
}


bool AudioStream::SendOggIdentPacket(IMediaSample* pSample)
{
    assert(pSample);
        
    const char* const id = m_pTrack->GetCodecId();
    id;
    assert(id);    
    assert(_stricmp(id, "A_VORBIS") == 0);
    
    const bytes_t& cp = m_pTrack->GetCodecPrivate();
    assert(!cp.empty());
    
    const ULONG cp_size = static_cast<ULONG>(cp.size());
    assert(cp_size > 0);
    
    const BYTE* const begin = &cp[0];
    const BYTE* const end = begin + cp_size;
    end;
    
    const BYTE* p = begin;
    assert(p < end);
    
    const BYTE n = *p++;
    n;
    assert(n == 2);
    assert(p < end);
    
    const ULONG id_len = *p++;  //TODO: don't assume < 255
    id_len;
    assert(id_len < 255);
    assert(id_len > 0);
    assert(p < end);
    
    const ULONG comment_len = *p++;  //TODO: don't assume < 255
    comment_len;
    assert(comment_len < 255);
    assert(comment_len > 0);
    assert(p < end);
    
    //p points to first header
    
    const BYTE* const id_hdr = p;
    
    const long size = pSample->GetSize();
    size;
    assert(size >= 0);
    assert(ULONG(size) >= id_len);
    
    BYTE* buf;
    
    HRESULT hr = pSample->GetPointer(&buf);
    assert(SUCCEEDED(hr));
    assert(buf);
    
    memcpy(buf, id_hdr, id_len);
    
    hr = pSample->SetActualDataLength(id_len);
    assert(SUCCEEDED(hr));

    hr = pSample->SetPreroll(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetMediaType(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetDiscontinuity(TRUE /* m_bDiscontinuity */ );  //TODO
    assert(SUCCEEDED(hr));

    //TODO    
    //set by caller:
    //m_bDiscontinuity = false;

    hr = pSample->SetMediaTime(0, 0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetSyncPoint(FALSE);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetTime(0, 0);
    assert(SUCCEEDED(hr));

    m_preroll = &AudioStream::SendOggCommentPacket;
    
    return true;  //preroll - don't send payload
}


bool AudioStream::SendOggCommentPacket(IMediaSample* pSample)
{
    assert(pSample);
        
    const char* const id = m_pTrack->GetCodecId();
    id;
    assert(id);    
    assert(_stricmp(id, "A_VORBIS") == 0);
    
    const bytes_t& cp = m_pTrack->GetCodecPrivate();
    assert(!cp.empty());
    
    const ULONG cp_size = static_cast<ULONG>(cp.size());
    assert(cp_size > 0);
    
    const BYTE* const begin = &cp[0];
    const BYTE* const end = begin + cp_size;
    end;
    
    const BYTE* p = begin;
    assert(p < end);
    
    const BYTE n = *p++;
    n;
    assert(n == 2);
    assert(p < end);
    
    const ULONG id_len = *p++;  //TODO: don't assume < 255
    assert(id_len < 255);
    assert(id_len > 0);
    assert(p < end);
    
    const ULONG comment_len = *p++;  //TODO: don't assume < 255
    assert(comment_len < 255);
    assert(comment_len > 0);
    assert(p < end);
    
    //p points to first header
    
    const BYTE* const id_hdr = p;
    id_hdr;
    
    const BYTE* const comment_hdr = id_hdr + id_len;
    comment_hdr;
    
    const long size = pSample->GetSize();
    size;
    assert(size >= 0);
    assert(ULONG(size) >= comment_len);
    
    BYTE* buf;
    
    HRESULT hr = pSample->GetPointer(&buf);
    assert(SUCCEEDED(hr));
    assert(buf);
    
    memcpy(buf, comment_hdr, comment_len);
    
    hr = pSample->SetActualDataLength(comment_len);
    assert(SUCCEEDED(hr));

    hr = pSample->SetPreroll(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetMediaType(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetDiscontinuity(TRUE /* m_bDiscontinuity */ );  //TODO
    assert(SUCCEEDED(hr));

    //TODO    
    //set by caller:
    //m_bDiscontinuity = false;

    hr = pSample->SetMediaTime(0, 0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetSyncPoint(FALSE);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetTime(0, 0);
    assert(SUCCEEDED(hr));

    m_preroll = &AudioStream::SendOggSetupPacket;
    
    return true;  //preroll - don't send payload
}


bool AudioStream::SendOggSetupPacket(IMediaSample* pSample)
{
    assert(pSample);
        
    const char* const id = m_pTrack->GetCodecId();
    id;
    assert(id);    
    assert(_stricmp(id, "A_VORBIS") == 0);
    
    const bytes_t& cp = m_pTrack->GetCodecPrivate();
    assert(!cp.empty());
    
    const ULONG cp_size = static_cast<ULONG>(cp.size());
    assert(cp_size > 0);
    
    const BYTE* const begin = &cp[0];
    const BYTE* const end = begin + cp_size;
    
    const BYTE* p = begin;
    assert(p < end);
    
    const BYTE n = *p++;
    n;
    assert(n == 2);
    assert(p < end);
    
    const ULONG id_len = *p++;  //TODO: don't assume < 255
    assert(id_len < 255);
    assert(id_len > 0);
    assert(p < end);
    
    const ULONG comment_len = *p++;  //TODO: don't assume < 255
    assert(comment_len < 255);
    assert(comment_len > 0);
    assert(p < end);
    
    //p points to first header
    
    const BYTE* const id_hdr = p;
    id_hdr;
    
    const BYTE* const comment_hdr = id_hdr + id_len;
    comment_hdr;
    
    const BYTE* const setup_hdr = comment_hdr + comment_len;
    setup_hdr;
    assert(setup_hdr < end);
    
    const ptrdiff_t setup_len_ = end - setup_hdr;
    assert(setup_len_ > 0);
    
    const ULONG setup_len = static_cast<ULONG>(setup_len_);

    const long size = pSample->GetSize();
    size;
    assert(size >= 0);
    assert(ULONG(size) >= setup_len);
    
    BYTE* buf;
    
    HRESULT hr = pSample->GetPointer(&buf);
    assert(SUCCEEDED(hr));
    assert(buf);
    
    memcpy(buf, setup_hdr, setup_len);
    
    hr = pSample->SetActualDataLength(setup_len);
    assert(SUCCEEDED(hr));

    hr = pSample->SetPreroll(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetMediaType(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetDiscontinuity(TRUE /* m_bDiscontinuity */ );  //TODO
    assert(SUCCEEDED(hr));

    //TODO    
    //set by caller:
    //m_bDiscontinuity = false;

    hr = pSample->SetMediaTime(0, 0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetSyncPoint(FALSE);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetTime(0, 0);
    assert(SUCCEEDED(hr));

    m_preroll = &AudioStream::DoNothing;
    
    return true;  //don't send payload
}


}  //end namespace MkvParser
